# -*- coding: utf-8 -*-
# @Time : 2022/3/30 14:57
# @Author : XyZeng
import time

import numpy as np
import cv2

def gray2color(gray_img, colormap):
    '''

    :param gray_img:  (h,w)  max_value=num_class
    :param colormap:  (num_class,3) # colormap的索引对应了灰度值
    :return: color_img: (h,w,3) colormap的索引与label中灰度值一一顺序对应；如colormap[0]的颜色对应gray_image==0
    '''
    color_img = colormap[gray_img]

    return color_img


# 512x1024  time=0.05s
def color2gray(image_rgb, rgb_mapping):
    gray_img = np.zeros(shape=(image_rgb.shape[0], image_rgb.shape[1]), dtype=np.uint8)
    _st_rgb2gray = time.time()
    for map_idx, rgb in enumerate(rgb_mapping):
        idx = np.where((image_rgb[..., 0] == rgb[0]) & (image_rgb[..., 1] == rgb[1]) & (image_rgb[..., 2] == rgb[2]))
        gray_img[idx] = map_idx
    print("rgb 2 gray time used: {}".format(time.time() - _st_rgb2gray))
    return gray_img


# slow 512x1024  time=0.2s
def color2gray_slow(color_img, colormap):
    '''

    :param color_img:
    :param colormap:
    :param num_classes:
    :return:
    '''
    t0 = time.time()
    gray_img = np.zeros(shape=(color_img.shape[0], color_img.shape[1]), dtype=np.uint8)

    #
    for i in range(colormap.shape[0]):
        index = np.where(np.all(color_img == colormap[i], axis=-1))  # np.all true false
        gray_img[index] = i

    t1 = time.time()
    time_cost = round(t1 - t0, 3)
    print(f"color2label  cost time {time_cost}")
    return gray_img
